from DrissionPage import ChromiumPage, ChromiumOptions

class DouyinVideoExtractor:
    """一个用于提取抖音视频信息的类。"""
    
    def __init__(self, headless=True):
        """
        初始化浏览器配置。
        :param headless: 是否使用无头模式。
        """
        self.co = ChromiumOptions()
        if headless:
            self.co.headless()
        self.browser = None
    
    def __enter__(self):
        """上下文管理器入口，用于自动打开浏览器。"""
        self.browser = ChromiumPage(self.co)
        return self
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口，用于自动关闭浏览器。"""
        if self.browser:
            self.browser.quit()
    
    def get_video_url(self, video_url: str) -> str | None:
        """
        获取单个抖音视频的播放地址。
        :param video_url: 抖音视频的分享链接。
        :return: 视频的直接播放地址或 None。
        """
        # 在每次请求新视频前，清除浏览器缓存
        self.browser.clear_cache()
        try:
            self.browser.listen.start("/aweme/detail")
            self.browser.get(video_url)
            packet = self.browser.listen.wait(timeout=10)
            
            if packet:
                data = packet.response.body.get("aweme_detail", {})
                if "video" in data:
                    video_info = data["video"]
                    url_list = video_info.get("bit_rate", [{}])[0].get("play_addr", {}).get("url_list")
                    if url_list:
                        return url_list[-1]
        except Exception as e:
            print(f"处理 {video_url} 时发生错误: {e}")
            return None
        finally:
            if hasattr(self.browser, 'listen'):
                self.browser.listen.clear()
        return None